package com.shujia.flink.sink

import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.{RichSinkFunction, SinkFunction}
import org.apache.flink.streaming.api.scala._

import java.sql.{Connection, DriverManager, PreparedStatement}

object Demo4SinkFunction {
  def main(args: Array[String]): Unit = {
    /**
     * 实时统计单词的数量，实时将统计结果写入到mysql
     *
     */
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    val linesDS: DataStream[String] = env.socketTextStream("master", 8888)

    //统计单词的数量
    val countDS: DataStream[(String, Int)] = linesDS
      .flatMap(_.split(","))
      .map((_, 1))
      .keyBy(_._1)
      .sum(1)

    //使用自定义sink
    countDS.addSink(new MysqlSink)

    env.execute()
  }

  /**
   * 自定义sink, 实现SinkFunction接口，重写invoke方法
   * SinkFunction
   * RichSinkFunction: 比SinkFunction多了open和close方法用于打开和关联链接
   *
   */
  class MysqlSink extends RichSinkFunction[(String, Int)] {

    var stat: PreparedStatement = _
    var con: Connection = _

    /**
     * open方法是每一个task中执行一次， 在invoke之前执行， 一般用于初始化的操作，比如建立数据库链接
     */

    override def open(parameters: Configuration): Unit = {
      println("执行open方法，创建数据库链接")
      //1、加载驱动
      Class.forName("com.mysql.jdbc.Driver")
      //2、建立链接
      con = DriverManager.getConnection("jdbc:mysql://master:3306/bigdata", "root", "123456")
      //3、编写sql
      stat = con.prepareStatement("replace into word_count(word,num) values(?,?)")
    }

    /**
     * close方法，每一个task中执行一次，在任务取消时执行，一般用于回收资源，如果任务一致运行，永远不会执行
     */
    override def close(): Unit = {
      //6、关闭链接
      stat.close()
      con.close()
    }

    /**
     * invoke方法每一条数据执行一次,在invoke方法中将数据保存到外部系统，比如mysql
     *
     * @param kv      ： 一行数据
     * @param context ： 上下文对象
     */
    override def invoke(kv: (String, Int), context: SinkFunction.Context): Unit = {
      println("执行invoke，插入数据")
      //4、设置字段的值
      val (word: String, num: Int) = kv
      stat.setString(1, word)
      stat.setInt(2, num)
      //5、执行插入
      stat.execute()
    }
  }
}
